import numpy as np
import os
import time
import subprocess
from mpi4py import MPI


def get_memory():
    cmd = "nvidia-smi --query-gpu=memory.used --format=csv | grep -v memory | awk '{print $1}'"
    output = subprocess.check_output(["bash", "-c", cmd])
    mem = [int(s.strip()) for s in output.decode().split('\n') if s.strip()] 
    return np.array(mem, dtype=np.int64)


def get_free_output():
    return subprocess.check_output(['free', '-h']).decode()


log_file = 'mem.log'
with open(log_file, 'w') as f:
    pass

comm = MPI.COMM_WORLD
cur_max_mem = None
while True: 
    max_mem = get_memory()
    gpu_num = max_mem.size
    max_mem_allgather = np.zeros([gpu_num * comm.size], dtype=np.int64) 
    comm.Allgather(max_mem, max_mem_allgather)
    max_mem_allgather = max_mem_allgather.reshape([comm.size, gpu_num])

    if cur_max_mem is None:
        cur_max_mem = max_mem_allgather
    else:
        for i in range(comm.size):
            for j in range(gpu_num):
                cur_max_mem[i][j] = max([cur_max_mem[i][j], max_mem_allgather[i][j]]) 

    s = '\t'.join([str(i) for i in max_mem_allgather.flatten()]) 
    if comm.rank == 0:
        with open(log_file, 'a+') as f:
            f.write(s + '\n')

            f.write('\t'.join(['Node-Rank'] + ['GPU-' + str(s) for s in range(gpu_num)]))
            f.write('\n')
            for i in range(comm.size): 
                f.write('\t'.join(['Node-{:0>5d}'.format(i)] + [str(s) for s in cur_max_mem[i]]))
                f.write('\n')
            f.write(get_free_output())
            f.write('\n')

    time.sleep(5)
